package org.data.collection.tree;

import cn.hutool.core.util.ObjectUtil;
import org.data.collection.tree.basic.BinaryTree;

import java.util.Comparator;
import java.util.Optional;
import java.util.concurrent.atomic.AtomicReference;

/**
 * 二叉排序树相关操作
 *
 * @author LiKun
 * @date 2021/11/22 19:20
 */
public class BinarySortTree<K, V> extends BinaryTree<K, V> {
    public BinarySortTree() {
        this.comparator = null;
    }

    public BinarySortTree(Comparator<K> comparator) {
        this.comparator = comparator;
    }

    @Override
    public V get(K key) {
        AtomicReference<V> reference = new AtomicReference<>();
        Optional.ofNullable(this.getEntry(key)).ifPresent(entry -> reference.set(entry.getValue()));
        return reference.get();
    }

    @Override
    public V setValue(K key, V value) {
        AtomicReference<V> reference = new AtomicReference<>();
        Optional.ofNullable(this.getEntry(key)).ifPresent(entry -> {
            reference.set(entry.getValue());
            entry.setValue(value);
        });

        return reference.get();
    }

    @Override
    public Entry<K, V> getEntry(K key) {
        Entry<K, V> result = null;
        if (ObjectUtil.isNotNull(root)) {
            final Entry<K, V> entry = this.getParentEntry(key, root);
            if (key.equals(entry.getKey())) {
                result = entry;
            }
        }
        return result;
    }

    @Override
    public boolean put(K key, V value) {
        final Entry<K, V> entry = new Entry<>(key, value);

        if (ObjectUtil.isNotNull(root) || !this.validAndSetRoot(entry)) {
            final Entry<K, V> parentEntry = this.getParentEntry(key, root);
            if (ObjectUtil.isNotNull(parentEntry)) {
                if (this.compare(key, parentEntry.getKey()) < 0) {
                    parentEntry.setLeftTree(entry);
                } else {
                    parentEntry.setRightTree(entry);
                }
            }
            size.incrementAndGet();
        }

        return true;
    }

    /*------------------------------------------------private methods-------------------------------------------------*/

    /**
     * 校验并设置 root 节点的值
     *
     * @param entry 节点数据
     *
     * @return 是否成功设置 root 节点的值
     */
    private boolean validAndSetRoot(Entry<K, V> entry) {
        boolean setRoot = false;
        if (ObjectUtil.isNull(root)) {
            synchronized (this) {
                if (ObjectUtil.isNull(root)) {
                    this.root = entry;
                    setRoot = true;
                }
            }
        }
        return setRoot;
    }

    private BinaryTree.Entry<K, V> getParentEntry(K key, Entry<K, V> entry) {
        Entry<K, V> result = entry;

        if (ObjectUtil.isNotNull(entry)) {
            int compare = this.compare(key, entry.getKey());

            if (compare < 0) {
                Entry<K, V> leftEntry;
                if (ObjectUtil.isNotNull(leftEntry = entry.getLeftTree())) {
                    result = this.getParentEntry(key, leftEntry);
                }
            } else if (compare > 0) {
                Entry<K, V> rightEntry;
                if (ObjectUtil.isNotNull(rightEntry = entry.getRightTree())) {
                    result = this.getParentEntry(key, rightEntry);
                }
            }
        }

        return result;
    }
}
